import natural from "natural";
import fs from "fs-extra";
import { getData } from "./getData.js";
import { sampleSize } from "lodash-es";
const list = await getData();

/** 训练集与测试集的比例 */
const rate = 10;
const special: { [string: string]: number } = {
    ts: 15,
    js: 15,
    java: 10,
    css: 15,
    tsx: 10,
    html: 10,
    go: 10,
    less: 10,
    cpp: 30,
    python: 30,
};
const trans = list.map((i) => {
    return {
        ...i,
        files: sampleSize(i.files, i.type in special ? special[i.type] : rate),
    };
});
var classifier = new natural.BayesClassifier();
trans.forEach((i) => {
    console.log("训练集数目");
    console.log(i.type, i.files.length);
    i.files.forEach((content) => {
        classifier.addDocument(content, i.type);
    });
});
classifier.save("./src/result/classifier.json", (err, callback) => {
    if (err) console.log(err);
    console.log("保存完成");
});
